{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Geodesics in Heat implementation\n",
    "This notebook implements Geodesics in Heat [[Crane et al. 2014]](https://arxiv.org/pdf/1204.6216.pdf) for triangle and tet meshes.\n",
    "\n",
    "Compare to the C++ implementation in `experiments/geodesic_heat/main.cc`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('..')\n",
    "import mesh, differential_operators, sparse_matrices, numpy as np\n",
    "from tri_mesh_viewer import TriMeshViewer as Viewer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "volMesh = mesh.Mesh('../../examples/meshes/3D_microstructure.msh', degree=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Choose whether to work with the tet mesh or its boundary triangle mesh.\n",
    "m = volMesh\n",
    "#m = volMesh.boundaryMesh()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Choose a timestep proportional to h^2 where h is the average edge length.\n",
    "# (As discussed in section 3.2.4 of the paper)\n",
    "c = 4 / np.sqrt(3)\n",
    "t = c *  m.volume / m.numElements()\n",
    "# Choose source vertex/vertices for computing distances\n",
    "sourceVertices = [0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have not yet bound the sparse matrix manipulation and solver functionality of MeshFEM, so we use scipy for now:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy, scipy.sparse, scipy.sparse.linalg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Backwards Euler time stepping for heat equation $\\frac{\\mathrm{d}}{\\mathrm{d}t} = \\bigtriangleup u, \\, \\, u|_\\gamma = 1 \\, \\forall t$:\n",
    "\\begin{align}\n",
    "                         \\frac{u_t - u_0}{t} &= \\bigtriangleup u_t \\\\\n",
    " \\Longrightarrow \\quad M \\frac{u_t - u_0}{t} &= -L u_t    \\quad \\text{(positive FEM Laplacian discretizes $-\\bigtriangleup$)} \\\\\n",
    " \\Longrightarrow \\quad \\underbrace{(M + t L)}_A u_t &= M u_0\n",
    " \\end{align}\n",
    "where $\\gamma$ is the domain from which we wish to compute distances (here given by `sourceVertices`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "L = differential_operators.laplacian(m).compressedColumn()\n",
    "M = differential_operators.mass(m, lumped=False).compressedColumn()\n",
    "A = L + t * M\n",
    "\n",
    "mask = np.ones(m.numVertices(), dtype=np.bool)\n",
    "mask[sourceVertices] = False\n",
    "\n",
    "A_ff = A[:,  mask][mask, :]\n",
    "A_fc = A[:, ~mask][mask, :]\n",
    "\n",
    "# Solve (M + t L) u = 0 with the constraint u[sourceVertices] = 1\n",
    "u = np.ones(m.numVertices())\n",
    "u[mask] = scipy.sparse.linalg.spsolve(A_ff, -A_fc @ np.ones(len(sourceVertices)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute the heat gradients\n",
    "g = differential_operators.gradient(m, u)\n",
    "# Normalize the gradients to get an approximate gradient of the distance field\n",
    "X = -g / np.linalg.norm(g, axis=1)[:, np.newaxis]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fit a scalar field's gradients to these normalized gradients $X$ by solving a Poisson equation:\n",
    "\n",
    "\\begin{align}\n",
    "- \\bigtriangleup \\phi = -\\nabla \\cdot X \\quad &\\text{in } \\Omega \\\\\n",
    "\\frac{\\mathrm{d} \\phi}{\\mathrm{d} {\\bf n}} = {\\bf n} \\cdot X \\quad &\\text{on } \\partial \\Omega \\\\\n",
    "\\phi = 0 \\quad &\\text{on } \\gamma\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "divX = differential_operators.divergence(m, X)\n",
    "L_ff = L[:, mask][mask, :]\n",
    "heatDist = np.zeros(m.numVertices())\n",
    "heatDist[mask] = scipy.sparse.linalg.spsolve(L_ff, divX[mask]) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualize the approximate distance field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "view = Viewer(m, scalarField=heatDist)\n",
    "view.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
